/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "eye_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_eye_custom.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void eye_custom(GM_ADDR x, GM_ADDR z, EyeCustomTilingData tiling);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 8;
    uint32_t tileNum = 8;
    size_t tilingSize = sizeof(EyeCustomTilingData);
    uint32_t num_rows = 2;
    uint32_t num_columns = 32;
    std::vector<uint32_t> batch_shape = {2, 1};
    uint32_t dtype;
    size_t batch_size = 1;
    for (const auto& dim : batch_shape) {
        batch_size *= dim;
    }
    // The byte size depends on dtype
    size_t data_size = sizeof(DTYPE_Y);
    // Check the DTYPE_Y in common.h and set dtype
    if (std::is_same<DTYPE_Y, _Float32>::value) {
        dtype = 0;
    } else if (std::is_same<DTYPE_Y, _Float16>::value) {
        dtype = 1;
    } else if (std::is_same<DTYPE_Y, int32_t>::value) {
        dtype = 2;
    } else if (std::is_same<DTYPE_Y, double>::value) {
        dtype = 3;
    }
    size_t inputByteSize = blockDim * batch_size * num_rows * num_columns * data_size;
    size_t outputByteSize = blockDim * batch_size * num_rows * num_columns * data_size;

#ifdef ASCENDC_CPU_DEBUG
    uint8_t *x = (uint8_t *)AscendC::GmAlloc(inputByteSize);
    uint8_t *z = (uint8_t *)AscendC::GmAlloc(outputByteSize);

    ReadFile("./input/input_x.bin", inputByteSize, x, inputByteSize);

    EyeCustomTilingData *tiling = (EyeCustomTilingData *)AscendC::GmAlloc(tilingSize);
    tiling->size = blockDim * batch_size * num_rows * num_columns;
    tiling->tileNum = tileNum;
    tiling->num_rows = num_rows;
    tiling->num_columns = num_columns;
    tiling->batch_shape_numel = batch_size;
    tiling->dtype = dtype;

    AscendC::SetKernelMode(KernelMode::AIV_MODE);

    ICPU_RUN_KF(eye_custom, blockDim, x, z,
                *reinterpret_cast<EyeCustomTilingData *>(tiling)); // use this macro for cpu debug

    WriteFile("./output/output_z.bin", z, outputByteSize);

    AscendC::GmFree((void *)x);
    AscendC::GmFree((void *)z);
    AscendC::GmFree((void *)tiling);
#else
    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    AddCustomTilingData *tiling;
    uint8_t *xHost, *yHost, *zHost;
    uint8_t *xDevice, *yDevice, *zDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);

    CHECK_ACL(aclrtMallocHost((void **)(&xHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&yHost), inputByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&zHost), outputByteSize));
    CHECK_ACL(aclrtMalloc((void **)&xDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&yDevice, inputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&zDevice, outputByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/input_x.bin", inputByteSize, xHost, inputByteSize);
    ReadFile("./input/input_y.bin", inputByteSize, yHost, inputByteSize);

    CHECK_ACL(aclrtMemcpy(xDevice, inputByteSize, xHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(yDevice, inputByteSize, yHost, inputByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    ACLRT_LAUNCH_KERNEL(add_custom)(blockDim, stream, xDevice, yDevice, zDevice, tiling);
    CHECK_ACL(aclrtSynchronizeStream(stream));

    CHECK_ACL(aclrtMemcpy(zHost, outputByteSize, zDevice, outputByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_z.bin", zHost, outputByteSize);

    CHECK_ACL(aclrtFree(xDevice));
    CHECK_ACL(aclrtFree(yDevice));
    CHECK_ACL(aclrtFree(zDevice));
    CHECK_ACL(aclrtFreeHost(xHost));
    CHECK_ACL(aclrtFreeHost(yHost));
    CHECK_ACL(aclrtFreeHost(zHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());
#endif
    return 0;
}
